#include "color.h"

float LabDistance2(const ColorLab &a, const ColorLab &b)
{
	float res = 0;
	for (int i = 0; i < 3; i++)
		res += (a[i] - b[i]) * (a[i] - b[i]);
	return res;
}

ColorLab Weight_Mul(const ColorLab &lab,float Weight)
{
	ColorLab res;
	for(int i = 0 ; i < 3 ; i++)
		res[i] = lab[i]*Weight;
	return res;
}

void LabSort(ColorLabList &colors)
{
	int l = colors.size();
	for (int i = l - 1; i > 0; i--)
		for (int j = 0; j < i; j++)
			if (colors[j][0] > colors[j + 1][0])
                std::swap(colors[j],colors[j+1]);
    return;
}

ColorLab RGB2Lab(int R, int G, int B)
{
	if(R==0 && G==0 && B==0)
		return ColorLab(0,0,0);
	float r = R*1.0f / 255.0f,
		  g = G*1.0f / 255.0f,
		  b = B*1.0f / 255.0f,
		  x, y, z;

	r = (r > 0.04045f) ? pow((r + 0.055f) / 1.055f, 2.4f) : r / 12.92f;
	g = (g > 0.04045f) ? pow((g + 0.055f) / 1.055f, 2.4f) : g / 12.92f;
	b = (b > 0.04045f) ? pow((b + 0.055f) / 1.055f, 2.4f) : b / 12.92f;

	x = (r * 0.4124f + g * 0.3576f + b * 0.1805f) / 0.95047f;
	y = (r * 0.2126f + g * 0.7152f + b * 0.0722f) / 1.00000f;
	z = (r * 0.0193f + g * 0.1192f + b * 0.9505f) / 1.08883f;

	x = (x > 0.008856f) ? pow(x, 1.0f/3.0f) : (7.787f * x) + 16.0f/116;
	y = (y > 0.008856f) ? pow(y, 1.0f/3.0f) : (7.787f * y) + 16.0f/116;
	z = (z > 0.008856f) ? pow(z, 1.0f/3.0f) : (7.787f * z) + 16.0f/116;

	ColorLab Lab;

	Lab[0] = (116.0f * y) - 16.0f;
	Lab[1] = 500.0f * (x - y);
	Lab[2] = 200.0f * (y - z);
	return Lab;
}

ColorRGBf Lab2RGBf(const ColorLab& Lab)
{
	float y = (Lab[0] + 16) / 116,
		  x = Lab[1] / 500 + y,
		  z = y - Lab[2] / 200,
		  r, g, b;

	x = 0.95047f * ((x * x * x > 0.008856f) ? x * x * x : (x - 16.0f/116) / 7.787f);
	y = 1.00000f * ((y * y * y > 0.008856f) ? y * y * y : (y - 16.0f/116) / 7.787f);
	z = 1.08883f * ((z * z * z > 0.008856f) ? z * z * z : (z - 16.0f/116) / 7.787f);

	r = x *  3.2406f + y * -1.5372f + z * -0.4986f;
	g = x * -0.9689f + y *  1.8758f + z *  0.0415f;
	b = x *  0.0557f + y * -0.2040f + z *  1.0570f;

	r = (r > 0.0031308f) ? (1.055f * pow(r, 1.0f/2.4f) - 0.055f) : 12.92f * r;
	g = (g > 0.0031308f) ? (1.055f * pow(g, 1.0f/2.4f) - 0.055f) : 12.92f * g;
	b = (b > 0.0031308f) ? (1.055f * pow(b, 1.0f/2.4f) - 0.055f) : 12.92f * b;

	ColorRGBf RGBf;
	RGBf[0] = r;
	RGBf[1] = g;
	RGBf[2] = b;
	return RGBf;
}

ColorRGB Lab2RGB(const ColorLab& Lab)
{
	float y = (Lab[0] + 16) / 116,
		  x = Lab[1] / 500 + y,
		  z = y - Lab[2] / 200,
		  r, g, b;

	x = 0.95047f * ((x * x * x > 0.008856f) ? x * x * x : (x - 16.0f/116) / 7.787f);
	y = 1.00000f * ((y * y * y > 0.008856f) ? y * y * y : (y - 16.0f/116) / 7.787f);
	z = 1.08883f * ((z * z * z > 0.008856f) ? z * z * z : (z - 16.0f/116) / 7.787f);

	r = x *  3.2406f + y * -1.5372f + z * -0.4986f;
	g = x * -0.9689f + y *  1.8758f + z *  0.0415f;
	b = x *  0.0557f + y * -0.2040f + z *  1.0570f;

	r = (r > 0.0031308f) ? (1.055f * pow(r, 1.0f/2.4f) - 0.055f) : 12.92f * r;
	g = (g > 0.0031308f) ? (1.055f * pow(g, 1.0f/2.4f) - 0.055f) : 12.92f * g;
	b = (b > 0.0031308f) ? (1.055f * pow(b, 1.0f/2.4f) - 0.055f) : 12.92f * b;

	ColorRGB RGB;
	RGB[0] = MAX(0, MIN(1, r)) * 255;
	RGB[1] = MAX(0, MIN(1, g)) * 255;
	RGB[2] = MAX(0, MIN(1, b)) * 255;
	return RGB;
}

void TransferColor(const Mat &inputimg, Mat &outputimg,const Palette & palette)
{
    vector<ColorL> L_change,L_origin;
    ColorLabList origin = palette.get_orinPalette();
    ColorLabList edited = palette.get_editPalette();
    L_change.push_back(0);
    L_origin.push_back(0);
    for(int i = 1 ; i < edited.size() ; i++)
    {
        L_change.push_back(edited[i][0]);
        L_origin.push_back(origin[i][0]);
    }
    L_change.push_back(100);

    L_origin.push_back(100);

    uchar* poutdata = outputimg.data;
	uchar* pindata = inputimg.data;
    int visable_palette_size = palette.get_K();
    Mat Lambda = palette.get_Lambda();
    float sigma = palette.get_Sigma();
    for (int i = 0; i < outputimg.rows; i ++)
    {
        for(int j = 0 ; j < outputimg.cols ; j++,poutdata+=3,pindata+=3)
        {
            int B=*pindata, G=*(pindata+1), R=*(pindata+2);
            ColorLab Lab = RGB2Lab(R,G,B);
            ColorLab res(0,0,0);
            float L_after = TransferLuminance(Lab[0],L_origin,L_change);
            //float tempsum = 0;

            for (int p = 0; p < visable_palette_size ; p++)
            {
                ColorLab v = TransferAB(origin[p+1], edited[p+1] , L_after , Lab);
                v[0] = L_after;
                float omega = Omega(origin, Lab, p, Lambda, sigma);
                v = Weight_Mul(v, omega);
                res = LabAdd(res, v);
               // tempsum+=omega;
            }

            //cout<<tempsum<<endl;
			res[0]=L_after;
            ColorRGB RGB = Lab2RGB(res);
            *(poutdata) = RGB[2];
            *(poutdata+1) = RGB[1];
            *(poutdata+2) = RGB[0];
        }
		if(i%100==0)
			cout<<i;
    }
    return;
}

ColorL TransferLuminance(ColorL L,const vector<ColorL> &L_origin,const vector<ColorL> &L_change)
{
    int i;
	for (i = 0; i < L_origin.size() -1 ; i++)
		if (L >= L_origin[i] && L <= L_origin[i + 1])
			break;
	ColorL l1 = L_origin[i];
	ColorL l2 = L_origin[i + 1];
	float k = (l1 == l2 ? 1 : (L - l1) / (l2 - l1));
	ColorL L1 = L_change[i];
	ColorL L2 = L_change[i + 1];
	ColorL L_res = (L2 - L1) * k + L1;
	return L_res;
}

ColorLab TransferAB(const ColorLab &input,const ColorLab &edited, ColorL L,const ColorLab &x)
{
	ColorLab C;
	C[0] = L, C[1] = input[1], C[2] = input[2];
	ColorLab Cp;
	Cp[0] = L, Cp[1] = edited[1], Cp[2] = edited[2];

	if (LabDistance2(C, Cp) < 0.0001) {
		return x;
	}
	ColorLab x0 = LabAdd(x, LabSub(Cp, C));
	ColorLab Cb = LabFindBoundary(C, Cp);
	ColorLab xb;
	if (isOutGamut(x0))
		xb = LabBoundary(Cp, x0);
    else
		xb = LabFindBoundary(x, x0);
    float dxx = LabDistance2(xb, x);
	float dcc = LabDistance2(Cb, C);
	float l2 = std::min(1.0f, dxx / dcc);
	ColorLab xbn = LabNorm(LabSub(xb, x));
	float x_x = sqrt(LabDistance2(C, Cp) * l2); //||x'-x||
	ColorLab xp = LabAdd(x, Weight_Mul(xbn, x_x));
	return xp;
}


ColorLab LabSub(const ColorLab &a, const ColorLab &b)
{
    ColorLab res;
    for(int i = 0 ; i < 3 ; i++)
    {
        res[i] = a[i]-b[i];
    }
    return res;
}

ColorLab LabAdd(const ColorLab &a, const ColorLab &b)
{
    ColorLab res;
    for(int i = 0 ; i < 3 ; i++)
    {
        res[i] = a[i]+b[i];
    }
    return res;
}

ColorLab LabNorm(const ColorLab &a)
{
	float d = sqrt(LabDistance2(a, ColorLab(0,0,0)));
	ColorLab res;
	for (int  i = 0; i < 3; i++) {
		res[i]=a[i] / d;
	}
	return res;
}

ColorLab LabFindBoundary(const ColorLab &p1,const ColorLab &p2)
 {
     if (isOutGamut(p2))
     {
         return LabBoundary(p1,p2);
     }
     else
    {
        return LabFindBoundary(p2,LabAdd(p2,LabSub(p2,p1)));
    }
}

ColorLab LabBoundary(const ColorLab& pin, const ColorLab& pout)
{
    ColorLab mid = Weight_Mul(LabAdd(pin , pout),0.5);
    if ( LabDistance2(pin,pout)<0.001 )
    {
        return mid;
    }
    if (isOutGamut(mid))
    {
        return LabBoundary(pin, mid);
    }
    else
    {
        return LabBoundary(mid, pout);
    }
}

float Omega(const ColorLabList &origin, const ColorLab &Lab, int i, const Mat & Lambda, float sigma)
{
	float sum = 0;
	for (int j = 0; j < origin.size()-1; j++) {
		sum += Lambda.at<float>(i,j) * Phi(LabDistance2(origin[j+1], Lab),sigma);
	}
	return sum;
}

float Phi(float r2, float sigma)
{
	return exp(- r2 / (2 * sigma * sigma));
}

bool isOutGamut(const ColorLab &a)
{
    ColorRGBf RGB = Lab2RGBf(a);
    for (int i = 0; i < 3; i++)
        if (RGB[i] < 0 || RGB[i] > 1)
				return true;

    return false;
}
